const builtin = @import("builtin");
const csr = @import("riscv").csr;
const root = @import("kernel.zig");
const trap = @import("trap.zig");

const print = root.print;

extern fn switch_to(to: *Context) callconv(.c) void; // 返回类型不能是 noreturn

const Context = packed struct {
    // Layout must be kept in sync with switch_to
    // ignore x0
    ra: usize = 0,
    sp: usize,
    gp: usize = 0,
    tp: usize = 0,
    t0: usize = 0,
    t1: usize = 0,
    t2: usize = 0,
    s0: usize = 0,
    s1: usize = 0,
    a0: usize = 0,
    a1: usize = 0,
    a2: usize = 0,
    a3: usize = 0,
    a4: usize = 0,
    a5: usize = 0,
    a6: usize = 0,
    a7: usize = 0,
    s2: usize = 0,
    s3: usize = 0,
    s4: usize = 0,
    s5: usize = 0,
    s6: usize = 0,
    s7: usize = 0,
    s8: usize = 0,
    s9: usize = 0,
    s10: usize = 0,
    s11: usize = 0,
    t3: usize = 0,
    t4: usize = 0,
    t5: usize = 0,
    t6: usize = 0,

    pub fn new(func: fn () void, thread_stack: []u8) Context {
        return .{
            .ra = @intFromPtr(&func),
            .sp = @intFromPtr(&thread_stack[thread_stack.len - 1]),
        };
    }
};

const Task = struct {
    // 栈太小的话，会覆盖 context 的空间，mepc 地址不对导致执行 mret 异常
    const STACK_SIZE = if (builtin.mode == .Debug) 2048 else 1024;

    context: Context,
    stack: [STACK_SIZE]u8 = @splat(65), // 这里的默认初始化也没有用到，two_tasks 是 undefined
};

// 为了让 switch_to 能访问到，它必须是全局变量
const MAX_TASKS = 2;
pub var two_tasks: [MAX_TASKS]Task = undefined; // 在堆里分配会好写一点吗？

var top: usize = 0;
var current: usize = 1;

pub fn info() void {
    for (&two_tasks, 0..) |*task, i| {
        print("task{d}: stack: {*} -> {*} {c}\n", .{
            i,
            &task.stack[0],
            &task.stack[Task.STACK_SIZE - 1],
            task.stack[0],
        });
    }
}

pub fn schedule() void { // 返回类型不能是 noreturn
    current = (current + 1) % top;
    switch_to(&two_tasks[current].context);
}

pub inline fn yield() void {
    schedule();
}

pub fn init() void {
    // 初始化 mscratch
    csr.mscratch.write(0);

    inline for (.{ user_task0, user_task1 }) |func| {
        // 没法用这种方法创建结构体的实例
        //two_tasks.tasks[two_tasks.top] = .create(func);

        // 下面的方法可以得到预期结果，先创建，再赋值
        // 但是中间变量会多复制一次。。
        //var task = Task{};
        //task.context = .new(func, &task.stack);
        //two_tasks[top] = task;

        two_tasks[top].context = .new(func, &two_tasks[top].stack);
        top += 1;
    }
}

fn user_task0() void {
    print("Task 0: Created!\n", .{});

    while (true) {
        print("Task 0: Running...\n", .{});
        trap.tests(); // exception!!
        delay(1000);
        yield();
    }
}

fn user_task1() void {
    print("Task 1: Created!\n", .{});

    while (true) {
        print("Task 1: Running...\n", .{});
        delay(1000);
        yield();
    }
}

fn delay(x: usize) void {
    var i: usize = 50000 * x;
    const ptr: *volatile usize = &i;
    while (ptr.* > 0) {
        ptr.* -= 1;
    }
}
